#处理所有与板块相关的基础函数


from struct import unpack
import pandas as pd
import os

# TDX_BLOCK_PATH = 'D:/new_tdx/T0002/hq_cache/'
TDX_BLOCK_PATH = './ycdata/blockfile/'

""" block_gn.dat,_fg.dat,_zs.dat  """
def zs_block(blk='zs'):
    #读取指数板块对应的二进制文件，得到各个指数所包含的各个个股信息
    bf = get_block_file(blk)
    # print('bf = ',bf)

    t = get_block_zs_yczs_loc(blk)
    t_list = t.values.tolist()
    # print('t_list = ',t_list)
    df = pd.merge(t,bf,on='name')
    return df

#通过读取概念板块的信息以及板块名和代码的对应关系，合并成完整的板块信息。
def gn_block(blk='gn') :
    bf = get_block_file(blk)
    # if blk == 'gn':
    #     print('bf . length = ',bf.__len__())
    #     print('bf = ',bf)
    # bf.drop(bf[bf['name'].isin(del_row[blk])].index,inplace=True)
    # bf['name'] = bf['name'].replace(mapping[blk],regex=True)

    t = get_block_zs_tdx_loc(blk)
    # t_list = t.values.tolist()
    # if blk == 'gn':
    #     print('t . length = ',t.__len__())
    #     print('t = ',t)
    #目前的tdxzscfg3中，并没有包含各类指数（例如沪深300）的对应的代码。t中并没有先关信息，所以直接返回bf，而不与t进行合并。
    if (blk == 'zs'):
        return bf
    # del t['block']

    df = pd.merge(t,bf,left_on='block',right_on='name')
    del df['block']
    del df['name_y']

    # if blk == 'gn':
    #     print('df . length = ',df.__len__())
    #     print('df = ',df)
    return df


def hy_block(blk='hy'):
    #begintime = datetime.datetime.now()
    stocklist = get_stock_hyblock_tdx_loc()
    #print(stocklist)
    blocklist = get_block_zs_tdx_loc(blk)
    #blocklist = blocklist.drop(blocklist[blocklist['name'].str.contains('TDX')].index)
    blocklist['block5'] = blocklist['block'].str[0:5]
    #print(blocklist)
    blocklist['num'] = 0
    blocklist['stocks'] = ''
    for i in range(len(blocklist)):
        blockkey = blocklist.iat[i, 2]
        if (len(blockkey) == 5):
            datai = stocklist[stocklist['block5'] == blockkey]  # 根据板块名称过滤
        else:
            datai = stocklist[stocklist['block'] == blockkey]  # 根据板块名称过滤
        # 板块内进行排序填序号
        datai = datai.sort_values(by=['code'], ascending=[True])
        #datai.reset_index(drop=True, inplace=True)
        codelist = datai['code'].tolist()

        blocklist.iat[i, 4] = len(codelist)
        blocklist.iat[i, 5] = str(codelist)
    blocklist = blocklist.drop(blocklist[blocklist['num'] == 0].index)
    #endtime = datetime.datetime.now()
    #print('Cost ' + str((endtime - begintime).seconds) + ' seconds')
    #print(blocklist)

    return blocklist

#返回各个板块对应的名称、代码对应关系。（只有风格、行业、风格等板块，没有指数对应的代码）
def get_block_zs_tdx_loc(block='hy'):

    buf_line = read_file_loc(TDX_BLOCK_PATH+'tdxzs3.cfg', '|')

    mapping = {'hy': '2', 'dq': '3', 'gn': '4', 'fg': '5', 'yjhy': '12', 'zs': '6'}
    df = pd.DataFrame(buf_line, columns=['name', 'code', 'type', 't1', 't2', 'block'])
    dg = df.groupby(by='type')
    # df.to_excel('block.xlsx')
    if (block == 'zs'):
        return df
    temp = dg.get_group(mapping[block]).reset_index(drop=True)
    temp.drop(temp.columns[[2, 3, 4]], axis=1, inplace=True)
    #temp.to_excel('tdxzs3.xlsx', index=False)
    return temp

#读取各个主要指数对应的code,这个文件是自己根据格式输入的。
def get_block_zs_yczs_loc(block='zs'):
    buf_line = read_file_loc_utf8(TDX_BLOCK_PATH+'tdxzs_yc.cfg', '|')
    df = pd.DataFrame(buf_line, columns=['name', 'code'])
    return df

#返回各个板块对应的名称以及所包含的个股数量以及各个个股
def get_block_file(block='gn'):
    file_name = f'block_{block}.dat'
    #print(PATH + file_name)
    with open(TDX_BLOCK_PATH + file_name, 'rb') as f:
        buff = f.read()

    head = unpack('<384sh', buff[:386])
    blk = buff[386:]
    blocks = [blk[i * 2813:(i + 1) * 2813] for i in range(head[1])]
    bk_list = []
    for bk in blocks:
        name = bk[:8].decode('gbk').strip('\x00')
        num, t = unpack('<2h', bk[9:13])
        stks = bk[13:(12 + 7 * num)].decode('gbk').split('\x00')
        bk_list = bk_list + [[name, block, num, stks]]
        # bk_list = bk_list + [[name, num, stks]]
    return pd.DataFrame(bk_list, columns=['name', 'tp', 'num', 'stocks'])
    # return pd.DataFrame(bk_list, columns=['name', 'num', 'stocks'])

#读取每个个股对应的交易市场、代码和板块之间的关系
def get_stock_hyblock_tdx_loc():
    buf_line = read_file_loc(TDX_BLOCK_PATH+'tdxhy.cfg', '|')
    buf_lis = []
    mapping = {'0': 'sz.', '1': 'sh.', '2': 'bj.'}
    for x in buf_line:
        # x[1] = mapping[x[0]] + x[1]
        buf_lis.append(x)

    df = pd.DataFrame(buf_lis, columns=['c0', 'code', 'block', 'c1', 'c2', 'c3'])
    # print(df)
    df.drop(df.columns[[0, 3, 4, 5]], axis=1, inplace=True)

    df = df[(df['block'] != '')]
    # df = df[df.code.str.startswith(('sz','sh'))]
    df['block5'] = df['block'].str[0:5]

    #df.to_excel('tdxhy.xlsx', index=False)
    return df

#读取文本文件，以|为分隔符
def read_file_loc(file_name, splits):
    with open(file_name, 'r') as f:
        buf_lis = f.read().split('\n')
    return [x.split(splits) for x in buf_lis[:-1]]

def read_file_loc_utf8(file_name, splits):
    with open(file_name, 'r',encoding='utf8') as f:
        buf_lis = f.read().split('\n')
    return [x.split(splits) for x in buf_lis[:-1]]